import idautils
import idc
import idaapi
import ida_name
import ida_lines
import ida_hexrays
import ida_struct 
import json
import re

sensitive_functions1 = [        
    'strcpy', 'memcpy', 'memmove',
    'strcat', 'sprintf', 'gets',
    'wcscpy', 'wcscat' , 'recv',
    'fgets', 'strncpy', 'memset',
    'strncat',
    ]

# 处理32位寄存器与64位寄存器之间的映射
REG_MAPPING = {
    "eax": "rax",
    "ebx": "rbx",
    "ecx": "rcx",
    "edx": "rdx",
    "esi": "rsi",
    "edi": "rdi",
    "esp": "rsp",
    "ebp": "rbp",
}

# 指令类型
ARITH_INSTRS = {
    "add": "add",
    "sub": "sub",
    "mul": "mul",  # 无符号乘法
    "imul": "imul",  # 有符号乘法
    "div": "div",  # 无符号除法
    "idiv": "idiv",  # 有符号除法
    "shl": "shl",  # 左移
    "shr": "shr",  # 右移
}

###############################################################################
# 工具函数
###############################################################################
def get_prev_instructions(start_ea, count=1):
    """
    从指令地址 start_ea 向前获取 count 条有效指令的地址列表（不含 start_ea 本身）。
    返回按照从近到远顺序排列的指令地址列表（最多 count 条）。
    """
    prev_addrs = []
    ea = start_ea
    while ea != idc.BADADDR and len(prev_addrs) < count:
        ea = idc.PrevHead(ea)
        if ea == idc.BADADDR:
            break
        prev_addrs.append(ea)
    return prev_addrs

def get_disasm(ea):
    """
    获取 ea 处的反汇编字符串。
    """
    return idc.GetDisasm(ea)

def get_func_start(ea):
    """
    获取 ea 所在函数的起始地址。
    """
    func = idaapi.get_func(ea)
    return func.start_ea if func else idc.BADADDR

def is_immediate(ea, opnd_idx):
    """
    判断当前指令第 opnd_idx 个操作数是否为立即数。
    """
    op_type = idc.GetOpType(ea, opnd_idx)
    return op_type == idc.o_imm

def handle_register(opnd):
    """
    检查操作数是否为 32 位寄存器，若是，则返回对应的 64 位寄存器。
    否则直接返回原寄存器。
    """
    # 将操作数转换为小写，避免大小写问题
    # opnd = opnd.lower()

    # 如果操作数是 32 位寄存器，映射为 64 位寄存器
    if opnd in REG_MAPPING:
        return REG_MAPPING[opnd]
    
    # 否则返回原始操作数
    return opnd

def get_stack_size(func_start, call_ea):
    stack_size = None

    # 遍历函数内所有指令（Heads 返回指定区域内的可反汇编地址）
    # 也可以用 idautils.FuncItems(func_start) 遍历函数内所有指令
    for ea in idautils.Heads(func_start, call_ea):
        # 如果已经到 strcpy 调用处就提前退出，不再继续往后找
        if ea >= call_ea:
            break

        # 获取当前指令的助记符
        mnem = idc.GetMnem(ea)
        # 如果是 sub，并且第一个操作数是 rsp
        if mnem == "sub":
            opnd0 = idc.GetOpnd(ea, 0)  # 第一个操作数
            # 处理寄存器
            opnd0 = handle_register(opnd0)
            if opnd0.lower() == "rsp":
                # 第二个操作数形如 0x30 或 30h，最好使用 GetOperandValue 来直接获取数值
                val = idc.GetOperandValue(ea, 1)
                if val is not None and val != 0:
                    stack_size = val
                    rsp_ea = ea
                    # 找到了就可以直接 break，也可以继续看后面是否有更大的栈空间
                    break
        elif mnem == "add":
            opnd0 = idc.GetOpnd(ea, 0)  # 第一个操作数
            # 处理寄存器
            opnd0 = handle_register(opnd0)
            if opnd0.lower() == "rsp":
                # 第二个操作数形如 0x30 或 30h，最好使用 GetOperandValue 来直接获取数值
                val = idc.GetOperandValue(ea, 1)
                if val is not None and val != 0:
                    stack_size = val - 0xFFFFFFFFFFFFFF00
                    rsp_ea = ea
                    # 找到了就可以直接 break，也可以继续看后面是否有更大的栈空间
                    break

    return stack_size, rsp_ea

def get_var_buffer(func_ea, var_name):
    """
    计算指定函数中某个堆栈变量的“栈大小”：
      如果不是最后一个变量，size = abs(offset - next_offset)
      如果是最后一个变量，size = abs(offset)
    并打印结果（十六进制与十进制）。
    """
    if func_ea == idc.BADADDR:
        print(f"❌ 未找到函数: {func_ea}")
        return None

    f = idaapi.get_func(func_ea)
    if not f:
        print(f"❌ 地址 0x{func_ea:X} 不是有效函数起始")
        return None

    # ---------- 2. 取得栈帧 ----------
    frame = idaapi.get_frame(f)
    if not frame or not frame.members:
        print("❌ 该函数没有栈帧或没有局部变量")
        return None

    members = []
    for m in frame.members:
        name = ida_struct.get_member_name(m.id)  # 🌟 只要一个参数
        off  = m.get_soff()                      # 负数，如 -0x1E
        members.append((abs(off), name))

    # 按离 rbp 的距离从大到小排序
    # members.sort(key=lambda x: x[0], reverse=True)
    members.sort(key=lambda x: x[0])

    idx = next((i for i, (_, n) in enumerate(members) if n == var_name), None)
    if idx is None:
        print(f"❌ 未找到变量 {var_name}")
        return None

    curr_abs = members[idx][0]

    if idx + 1 < len(members):
        next_abs = members[idx + 1][0]
        size = curr_abs - next_abs
    else:
        size = curr_abs  # 已是最后一个变量

    return size

def get_reg_chain(func_start, call_ea, reg_name):
    # 用于存储回溯分析时的区间信息
    # 例如，每找到一次 mov 指令，就记录 (startEA, endEA, 寄存器)
    traced_ranges = []
    func_start =  get_func_start(func_start)
    
     # 2. 初始化回溯的起点信息
    cur_ea  = call_ea      # 当前起点
    cur_reg = reg_name     # 当前追踪的寄存器

    # 3. 从 call_ea 开始，一直往上找到 func_start
    ea = call_ea
    while ea >= func_start:
        mnem  = idc.GetMnem(ea)         # 指令助记符
        opnd0 = idc.GetOpnd(ea, 0)      # 第一个操作数
        opnd1 = idc.GetOpnd(ea, 1)      # 第二个操作数
        opnd2 = idc.GetOpnd(ea, 2)
        
        # 处理寄存器
        opnd0 = handle_register(opnd0)
        opnd1 = handle_register(opnd1)
        
        # if(opnd1 == "rdx"):
        #     print(mnem)
        #     if is_immediate(ea, 1):
        #         print(idc.GetOperandValue(ea, 1))

        # 4. 如果是 mov 指令，并且目标寄存器 == 当前追踪寄存器
        if mnem == "mov" and opnd0.lower() == cur_reg.lower():
            if is_immediate(ea, 1):
                return idc.GetOperandValue(ea, 1), ea
                
            else :
                # 记录下此时的区间：(本条指令地址 ~ 当前EA, 以及寄存器)
                traced_ranges.append((ea, cur_ea, cur_reg))

                # 将当前寄存器更新为源操作数
                cur_reg = opnd1
                # 将“区间结束地址”更新为当前指令地址
                cur_ea = ea

                # 如果更新后发现源操作数就是 rsp，就退出回溯
                if opnd1.lower() == "rsp":
                    traced_ranges.append((func_start, cur_ea, "rsp"))
                    break
                
        elif mnem == "lea" and opnd0.lower() == cur_reg.lower():
            print("    op_type = {} , opnd0 = {}, opnd1 = {}, opnd2 = {}".format(mnem, opnd0, opnd1, opnd2))
                # 通用匹配 [寄存器+标识符]，标识符可以是 buf、var_4、myVar123 等
            match_var = re.match(r'\[(?P<reg>[A-Za-z0-9]+)\+(?P<var>[A-Za-z_][A-Za-z0-9_]*)\]', opnd1)
            if match_var:
                reg      = match_var.group('reg')   # 寄存器部分，如 "rbp"
                var_name = match_var.group('var')   # 标识符部分，如 "buf"、"var_4"
                print(f"    detect register+symbol: reg={reg}, var={var_name}")
                var_size = -get_var_buffer(func_start, var_name)
                print(f"    变量 {var_name} 的大小: 0x{var_size:X} ({var_size} 字节)")
                return  var_size, ea
                # 在这里对 reg 和 var_name 做你自己的处理，比如：
                #   - 查表：symbol_table[var_name]
                #   - 计算偏移：get_symbol_size(var_name)
            # 然后返回或赋值
            
            match = re.match(r'\[(?P<reg>[a-zA-Z0-9]+)(\+(?P<imm>-?\d+))?\]', opnd1)
            if match:
                reg = match.group('reg')  # 寄存器部分
                imm = match.group('imm')  # 立即数部分，可能为空
                # print(reg)
                reg_size, start_ea = get_reg_chain(func_start, ea, reg)
                reg_size = calculate_buffer(start_ea, ea, reg, reg_size, calculated_registers)
                # print(reg, imm, reg_size)
                return  reg_size + int(imm), ea
            elif opnd1.isalpha():  # 如果是纯字母
                for i, ch in enumerate(opnd1):
                    if ch.isupper() and i > 0:
                        return opnd1[i:]  # 返回从第一个大写字母开始的部分
                return opnd1.upper()  # 如果没有小写字母，返回大写字符串
            
        
        elif mnem == "call" and cur_reg.lower() == "rax":
            if opnd0 == "_strlen":
                # print(opnd0)
                reg_size = get_reg_chain(func_start, idc.PrevHead(ea), cur_reg)
                print("    op_type = {} , opnd0 = {}, opnd1 = {}, opnd2 = {}, reg_size = {}".format(mnem, opnd0, opnd1, opnd2, len(reg_size)))
                return len(reg_size), ea
            
        # 移动到上一条指令
        prev_ea = idc.PrevHead(ea)
        if prev_ea == idc.BADADDR or prev_ea < func_start:
            break
        ea = prev_ea
    
    return traced_ranges

def calculate_buffer(start_ea, end_ea, reg_name, buffer_size, calculated_registers):
    #TODO: MEM的建模
    """
    从 start_ea 到 end_ea，逐行分析，若遇到加减乘除、左移右移等指令，
    并且目标寄存器是 reg_name，则对 buffer_size 做相应运算。

    - 若另一个操作数是立即数，直接计算；
    - 若是另一个寄存器，则向上寻找它的 mov 初始化（mov regX, imm），
      并将 imm 作为本次运算操作数（示例中仅演示简化做法）。
    - 最终返回计算完成的 buffer_size。
    """
    ea = start_ea

    # 如果寄存器已经计算过并且当前计算范围的起始地址在某个终止位置之后，则直接继承之前的值
    if reg_name in calculated_registers:
        for last_end_ea, last_buffer_size in calculated_registers[reg_name]:
            if last_end_ea >= start_ea:
                buffer_size = last_buffer_size
                ea = last_end_ea
                break

    while ea < end_ea:
        mnem = idc.GetMnem(ea).lower()
        if mnem in ARITH_INSTRS:
            opnd0 = idc.GetOpnd(ea, 0)  # 目标寄存器
            opnd1 = idc.GetOpnd(ea, 1)  # 操作数
            opnd2 = idc.GetOpnd(ea, 2)  # 操作数
            
            # 如果目标寄存器是我们关心的寄存器
            if opnd0.lower() == reg_name.lower():
                op_type = mnem  # 获取操作符类型，如 "add"、"sub"、"mul" 等
                
                # 输出格式化字符串，判断哪些操作数不为空再输出
                output_str = "    op_type = {} ,".format(op_type)
                # 判断并输出每个操作数，如果不为空
                if opnd0:
                    output_str += " opnd0 = {}".format(opnd0)
                if opnd1:
                    if opnd0:  # 如果 opnd0 已经被输出过，添加逗号
                        output_str += ","
                    output_str += " opnd1 = {}".format(opnd1)
                if opnd2:
                    if opnd1 or opnd0:  # 如果 opnd1 或 opnd0 已经被输出过，添加逗号
                        output_str += ","
                    output_str += " opnd2 = {}".format(opnd2)
                print(output_str)
                
                # 检查第二个操作数是否为立即数
                val = None
                if is_immediate(ea, 1):
                    val = idc.GetOperandValue(ea, 1)
                    # print("    val = {}".format(val))
                elif opnd1:
                    # 如果是另一个寄存器，则向上寻找其初始化值
                    print(start_ea, ea, opnd1)
                    initial_size, reg_ea = get_reg_chain(start_ea, ea, opnd1)
                    print("    reg_ea = {}, initial_size = {} , reg = {}".format(hex(reg_ea), initial_size, opnd1))
                    # print(idc.PrevHead(ea))
                    val = calculate_buffer(reg_ea, ea, opnd1, initial_size, calculated_registers)
                    print("    val = {}, reg = {}".format(val, opnd1))
                elif is_immediate(ea, 2):
                    val = idc.GetOperandValue(ea, 2)
                elif opnd2:
                    # 如果是另一个寄存器，则向上寻找其初始化值
                    initial_size, reg_ea = get_reg_chain(start_ea, ea, opnd2)
                    print("    reg_ea = {}, initial_size = {} , reg = {}".format(hex(reg_ea), initial_size, opnd2))
                    # print(idc.PrevHead(ea))
                    val = calculate_buffer(reg_ea, ea, opnd2, initial_size, calculated_registers)
                    print("    val = {}, reg = {}".format(val, opnd2))
                                  
                if val is not None:
                    if op_type == "add":
                        if opnd0 == "rsp":
                            buffer_size = buffer_size + val - 0xFFFFFFFFFFFFFF00
                        else:
                            buffer_size += val
                    elif op_type == "sub":
                        buffer_size -= val
                    elif op_type == "mul":
                        buffer_size *= val
                    elif op_type == "shl":
                        buffer_size <<= val
                    elif op_type == "shr":
                        buffer_size >>= val        
                    elif op_type == "imul":
                        # print(reg_name.lower())
                        # 处理 imul 指令（双操作数形式）
                        if opnd1:  # 存在第二个操作数
                            buffer_size = val * idc.GetOperandValue(ea, 2) 
                        # 处理 imul 双操作数形式
                        else:
                            buffer_size *= val
                            # print("    val = {}, buffer_size = {}".format(val, buffer_size))
                            
            elif mnem == ("div" or "idiv") and reg_name.lower() == "rax":
                # 处理 idiv 指令（除法）
                if is_immediate(ea, 1):  # 除数是立即数
                    divisor = idc.GetOperandValue(ea, 1)
                    print("    divisor = {}".format(divisor))
                    if divisor != 0:
                        quotient = buffer_size // divisor
                        remainder = buffer_size % divisor
                        buffer_size = quotient  # 商存储在 rax
                        
                        #可以在这里处理余数存储到 rdx
                        if "rdx" not in calculated_registers:
                            calculated_registers["rdx"] = []

                        calculated_registers["rdx"].append((end_ea, remainder))
                        
                else:  # 除数是寄存器
                    initial_size, reg_ea = get_reg_chain(start_ea, ea, opnd1)
                    print("    reg_ea = {}, initial_size = {} , reg = {}".format(hex(reg_ea), initial_size, opnd1))
                    divisor = calculate_buffer(reg_ea, idc.PrevHead(ea), opnd1, initial_size, calculated_registers)
                    print("    divisor = {}, reg = {}, buffer_size = {}".format(divisor, reg_name, buffer_size))
                    quotient = buffer_size // divisor
                    remainder = buffer_size % divisor
                    buffer_size = quotient  # 商存储在 rax
                    
                     #可以在这里处理余数存储到 rdx
                    if "rdx" not in calculated_registers:
                        calculated_registers["rdx"] = []

                    calculated_registers["rdx"].append((end_ea, remainder))
        
        ea = idc.NextHead(ea)  # 获取下一条指令
    
    # 更新寄存器的计算结果，存储多个终止位置的结果
    if reg_name not in calculated_registers:
        calculated_registers[reg_name] = []
    
    calculated_registers[reg_name].append((end_ea, buffer_size))

    return buffer_size
    
calculated_registers = {}

def get_buffer(func_start, call_ea, reg_name):
    cur_ea  = call_ea      # 当前起点
    cur_reg = reg_name     # 当前追踪的寄存器

    # 3. 从 call_ea 开始，一直往上找到 func_start
    ea = call_ea
    while ea >= func_start:
        mnem  = idc.GetMnem(ea)         # 指令助记符
        opnd0 = idc.GetOpnd(ea, 0)      # 第一个操作数
        opnd1 = idc.GetOpnd(ea, 1)      # 第二个操作数
        opnd2 = idc.GetOpnd(ea, 2)
        
        if mnem == "mov" and opnd0.lower() == cur_reg.lower():
            if is_immediate(ea, 1):
                return idc.GetOperandValue(ea, 1)
            # 如果更新后发现源操作数就是 rsp，就退出回溯
            elif opnd1.lower() == "rsp":
                break
            else:
                get_buffer(func_start, ea, opnd1)
                
        elif mnem == "lea" and opnd0.lower() == cur_reg.lower():
            print("    op_type = {} , opnd0 = {}, opnd1 = {}, opnd2 = {}".format(mnem, opnd0, opnd1, opnd2))
                # 通用匹配 [寄存器+标识符]，标识符可以是 buf、var_4、myVar123 等
            match_var = re.match(r'\[(?P<reg>[A-Za-z0-9]+)\+(?P<var>[A-Za-z_][A-Za-z0-9_]*)\]', opnd1)
            if match_var:
                reg      = match_var.group('reg')   # 寄存器部分，如 "rbp"
                var_name = match_var.group('var')   # 标识符部分，如 "buf"、"var_4"
                print(f"    detect register+symbol: reg={reg}, var={var_name}")
                var_size = -get_var_buffer(func_start, var_name)
                print(f"    变量 {var_name} 的大小: 0x{var_size:X} ({var_size} 字节)")
                return  var_size, ea
                # 在这里对 reg 和 var_name 做你自己的处理，比如：
                #   - 查表：symbol_table[var_name]
                #   - 计算偏移：get_symbol_size(var_name)
            # 然后返回或赋值
            
            match = re.match(r'\[(?P<reg>[a-zA-Z0-9]+)(\+(?P<imm>-?\d+))?\]', opnd1)
            if match:
                reg = match.group('reg')  # 寄存器部分
                imm = match.group('imm')  # 立即数部分，可能为空
                # print(reg)
                reg_size, start_ea = get_reg_chain(func_start, ea, reg)
                reg_size = calculate_buffer(start_ea, ea, reg, reg_size, calculated_registers)
                # print(reg, imm, reg_size)
                return  reg_size + int(imm), ea
            elif opnd1.isalpha():  # 如果是纯字母
                for i, ch in enumerate(opnd1):
                    if ch.isupper() and i > 0:
                        return opnd1[i:]  # 返回从第一个大写字母开始的部分
                return opnd1.upper()  # 如果没有小写字母，返回大写字符串
            
        
        elif mnem == "call" and cur_reg.lower() == "rax":
            if opnd0 == "_strlen":
                # print(opnd0)
                reg_size = get_reg_chain(func_start, idc.PrevHead(ea), cur_reg)
                print("    op_type = {} , opnd0 = {}, opnd1 = {}, opnd2 = {}, reg_size = {}".format(mnem, opnd0, opnd1, opnd2, len(reg_size)))
                return len(reg_size), ea
            
        # 移动到上一条指令
        prev_ea = idc.PrevHead(ea)
        if prev_ea == idc.BADADDR or prev_ea < func_start:
            break
        ea = prev_ea
    

###############################################################################
# 缓冲区大小分析逻辑（示例化）
###############################################################################
def analyze_buffer_size_for_register(func_start, call_ea, reg_name):
    """
    从调用指令 call_ea 开始，向上回溯 reg_name 的来源，直到函数起始位置或找到最终赋值为止。
    并根据 mov / sub / shr 等指令，估算（示例化）缓冲区大小。
    
    返回一个整数，表示粗略估算的缓冲区大小；若无法确定，返回 None。
    """
    buffer_size = None
    stack_size = None
    
    #首先获取初始化变量的空间
    stack_size, rsp_ea = get_stack_size(func_start, call_ea)
    
    print("[+] Function variable stack size = 0x{:X}".format(stack_size if stack_size else 0))
    
    # get_buffer(func_start, call_ea, reg_name)
    
    traced_ranges = get_reg_chain(func_start, call_ea, reg_name)

    # 打印一下回溯得到的信息
    print("[+] Trace result (backward from 0x{:X}):".format(call_ea))
    for (start_ea, end_ea, reg) in traced_ranges:
        print("    0x{:X} -> 0x{:X}, reg = {}".format(start_ea, end_ea, reg))
        
    # 用来存储寄存器的计算结果，避免重复计算

    
    """
    从 traced_ranges 中处理每一项，通过 calculate_buffer 来更新 buffer_size。
    处理时从最后一条开始，并将每次的 buffer_size 传递给下一次计算。
    """
    
    register_offset = None
    
    # 从最后一条开始，首先处理 rsp_ea 到 start_ea 的 buffer_size
    if traced_ranges:
        start_ea, end_ea, reg = traced_ranges[-1]  # 获取最后一条的起始地址
        start_ea = rsp_ea
        register_offset = calculate_buffer(start_ea, end_ea, reg, 0, calculated_registers)
        print("    0x{:X} -> 0x{:X}, reg = {}, register_offset = {}".format(start_ea, end_ea, reg, hex(register_offset)))
        # print(buffer_size)

        # 对 traced_ranges 从后向前计算 buffer_size
        for i in range(len(traced_ranges) - 2, -1, -1):
            start_ea, end_ea, reg = traced_ranges[i]
            register_offset = calculate_buffer(start_ea, end_ea, reg, register_offset, calculated_registers)
            print("    0x{:X} -> 0x{:X}, reg = {}, register_offset = {}".format(start_ea, end_ea, reg, hex(register_offset))) 
    
    buffer_size = - register_offset - stack_size
    print("    0x{:X} -> 0x{:X}, reg = {}, buffer_size = {}".format(traced_ranges[-1][0], traced_ranges[0][1], reg, hex(buffer_size)))

    return buffer_size

def extract_variables_from_call(call_str):
    """
    使用正则表达式提取函数调用中的参数。
    示例输入：sscanf(a2, "%[^.].%[^.].%[^.].%s", v47, v46, s, v44)
    示例输出：['a2', 'v47', 'v46', 's', 'v44']
    """
    # 修改后的正则表达式，捕获括号内的内容，允许嵌套括号
    match = re.search(r'\((.*)\)', call_str)
    if match:
        params_str = match.group(1)
        # 通过逗号分隔参数，并去除空格
        params = [param.strip() for param in re.split(r',\s*(?![^(]*\))', params_str) if param.strip()]
        return params
    return []

def get_variable_type_and_size(ea, variables):
    """
    查找并返回指定变量的类型和大小。
    例如：变量 v46 的类型为 char v46[16]，返回 'v46:16'。
    """
    func = idaapi.get_func(ea)
    if not func:
        print(f"No function found at address 0x{ea:X}")
        return {}

    # 反编译函数
    cfunc = idaapi.decompile(func)
    if not cfunc:
        print(f"Failed to decompile function at address 0x{ea:X}")
        return {}
    
    # 假设 cfunc 是您想要处理的对象
    cfunc_str = str(cfunc)

    # 正则表达式匹配从 "{" 到第一个空行之间的部分
    match = re.search(r'\{(.*?)\n\s*\n', cfunc_str, re.DOTALL)

    # 获取匹配的内容
    if match:
        filtered_str = match.group(1)  # 删除多余的空格
    else:
        filtered_str = ""
        
    lines = filtered_str.strip().split("\n")
    
    variable_info = {}
    var_num = 1

    # 搜索函数头部的声明（通常在反编译代码的顶部）
    for var in variables:
        print(var)
        if '"' in var:  # 如果变量是一个字符串常量
            var_type = 'string'
            string_size = len(re.search(r'"(.*)"', var).group(1))  # 获取字符串长度
            var_size = string_size
                # 处理数字类型
        elif re.search(r'0[xX][0-9a-fA-F]+[uUlL]*', var):  # 如果变量是一个数字常量
            var_type = 'number'
            # 提取数字并转换为十进制
            num_match = re.search(r'0[xX]([0-9a-fA-F]+)[uUlL]*', var)
            if num_match:
                var_size = int(num_match.group(1), 16)  # 转换为十进制
            else:
                var_size = ''
        elif re.search(r'^\d+[uUlL]*$', var):  # 十进制常量
            var_type = 'number'
            # 提取数字部分并按十进制转换
            num_match = re.search(r'^(\d+)[uUlL]*$', var)
            if num_match:
                var_size = int(num_match.group(1), 10)
            else:
                var_size = ''
        # m_ptr = re.match(r'^&(\w+)$', var)
        elif re.search(r'^&(\w+)$', var):
            var_match = re.search(r'^&(\w+)$', var)
            var_type = 'pointer'
            var_size = '*'
            var = var_match.group(1)  # 去除 &，保留变量名
        # 2. 如果是 sizeof(s) 这种 sizeof + 变量
        elif var.startswith('sizeof(') and var.endswith(')'):
            var_type = 'num'
            inner = var[len('sizeof('):-1]  # 提取括号内的变量名
            var_size = inner 
            # var_size = variable_info[inner]['var_size']  # 从 var_sizes 字典中获取对应大小
        # 3. 如果是 (const char *)s 这种类型的指针
        elif var.startswith('(const char *)'):
            var_type = 'pointer'
            var_size = '*'
            var = var[len('(const char *)'):]  # 去除类型前缀，保留变量名
        else:
            for line in lines:
            # for var in variables:
                # print(var)
                if var in line:           
                    # 使用正则表达式匹配类型和大小
                    match = re.search(r'(\w+)\s+(\*?)' + re.escape(var) + r'(\[([0-9]+)\])?', line)
                    if match:
                        var_type = match.group(1)  # 获取类型，如 char
                        pointer = match.group(2)   # 获取是否为指针
                        var_size = match.group(3)  # 获取数组部分

                        if pointer:
                            # 如果是指针
                            var_type = 'pointer'
                            var_size = '*'
                        elif var_size:
                            # 如果是数组，获取大小
                            var_type = 'array'
                            size_match = re.search(r'\[([0-9]+)\]', var_size)
                            if size_match:
                                var_size = size_match.group(1)
                        else:
                            # 如果既不是指针也不是数组
                            var_size = ''

        variable_info[var] = {
            'var_num': var_num,  # 记录排名
            'var_type': var_type,
            'var_size': var_size
        }
                    
        var_num += 1
                    
    return variable_info

def get_decompiled_line_from_address(ea):
    class CallFinder(idaapi.ctree_visitor_t):
        def __init__(self, ea):
            idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST)
            self.ea = ea
            self.call_item = None
            self.found = False

        def visit_expr(self, expr):
            if expr.op == idaapi.cot_call and expr.ea == self.ea:
                self.call_item = expr
                self.found = True
                return 1
            return 0

    # 初始化 Hex-Rays 反编译插件
    if not ida_hexrays.init_hexrays_plugin():
        print("Hex-Rays decompiler is not available.")
        return ""

    # 获取指定地址的函数
    func = idaapi.get_func(ea)
    if not func:
        print(f"No function found at address 0x{ea:X}")
        return ""

    # 反编译函数
    cfunc = idaapi.decompile(func)
    if not cfunc:
        print(f"Failed to decompile function at address 0x{ea:X}")
        return ""

    # 使用 CallFinder 查找特定调用
    finder = CallFinder(ea)
    finder.apply_to(cfunc.body, None)

    if finder.found and finder.call_item:
        try:
            # 获取调用表达式
            call_item = finder.call_item

            # 获取参数信息
            args = []
            if call_item.a:
                for i in range(call_item.a.size()):
                    arg = call_item.a[i]
                    args.append(str(arg))

            # 获取反编译的代码行
            item_line = ida_lines.tag_remove(call_item.print1(None)).strip()
            if item_line:
                return item_line

            # 如果上面的方法失败，尝试遍历所有代码行
            sv = cfunc.get_pseudocode()
            for i in range(len(sv)):
                line_item = sv[i]
                if ea in [item.ea for item in cfunc.treeitems]:
                    return ida_lines.tag_remove(line_item.line).strip()

        except Exception as e:
            print(f"Error processing call at 0x{ea:X}: {str(e)}")

    # 如果无法找到，回退到原有模式
    match = re.search(r'call\s+([\w@]+)', str(ea))
    if match:
        called_func = match.group(1).lstrip('_')
        sv = cfunc.get_pseudocode()
        if sv:
            code_lines = [ida_lines.tag_remove(sv[i].line) for i in range(sv.size())]
            for code_line in code_lines:
                if re.search(r'\b' + re.escape(called_func) + r'\b', code_line):
                    return code_line.strip()

    return ""

def analyze_pointer_size(address, info):
    call_disasm = get_disasm(address).lower()
    
    print(call_disasm)
    if "call" not in call_disasm:
        print("\n不是调用位置")

    # 获取该调用所在的函数
    func_start = get_func_start(address)
    if func_start == idc.BADADDR:
        print("\n获取函数起始位置出错")

    print("\n[+] 在函数 0x{:x} 中发现调用 strcpy 的指令：0x{:x}".format(func_start, address))
    
    # 用于存储每条指令的第二个操作数
    operands = []
    operand_address = []

    # 3. 从调用 strcpy 的位置向上获取两行代码
    prev_addrs = get_prev_instructions(address, count=2)
    # 简单展示一下
    for ea in prev_addrs:
        print("    上一条指令: 0x{:x} : {}".format(ea, get_disasm(ea)))
            # GetOpnd(ea, 1) 表示取第二个操作数（索引从0开始）
        second_operand = idc.GetOpnd(ea, 1)
        operand_address.append(ea)
        operands.append(second_operand)
    
    pointer_size = analyze_buffer_size_for_register(func_start, operand_address[info['var_num'] - 1], operands[info['var_num'] - 1])
    
    info['var_size'] = pointer_size

# 读取JSON文件中的函数名和调用地址
# def read_sensitive_functions_from_json(json_file_path):
#     """
#     从指定的JSON文件读取敏感函数名和调用地址。
#     返回的字典包含'name'和'call_address'。
#     """
#     with open(json_file_path, 'r') as f:
#         data = json.load(f)
    
#     sensitive_funcs = []
#     if "direct_sensitive_func" in data:
#         for func in data["direct_sensitive_func"]:
#             name = func.get("name")
#             call_address = func.get("call_address")
#             if name and call_address:
#                 sensitive_funcs.append({"name": name, "call_address": call_address})
#     if "inline_sensitive_func" in data:
#         for func in data["inline_sensitive_func"]:
#             name = func.get("inline_func_name_predict")
#             call_address = func.get("call_address")
#             if name and call_address:
#                 sensitive_funcs.append({"name": name, "call_address": call_address})
    
#     return sensitive_funcs

def read_sensitive_functions_from_json(json_file_path):
    with open(json_file_path, 'r', encoding='utf-8') as f:
        return json.load(f)

def write_sensitive_functions_to_json(json_file_path, data):
    with open(json_file_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=4, ensure_ascii=False)

class SensitiveFunctionsHandler:
    def __init__(self, variable_info):
        self.variable_info = variable_info

    def handle_memcpy(self):
        """
        处理 memcpy 函数的特定逻辑，检查是否存在缓冲区溢出。
        """
        var_sizes = self._get_var_sizes([1, 2, 3])

        if var_sizes.get(1) is None or var_sizes.get(2) is None or var_sizes.get(3) is None:
            print("无法获取所有必要的变量信息")
            return False

        var_num_1_size = var_sizes[1]
        var_num_2_size = var_sizes[2]
        var_num_3_size = var_sizes[3]

        # 判断是否存在缓冲区溢出
        if var_num_3_size > var_num_1_size:
            return True
        return False

    def handle_strcpy(self):
        """
        处理 strcpy 函数的特定逻辑。
        """
        # 这里可以加入特定的判断逻辑，比如判断源缓冲区大小是否大于目标缓冲区
        var_sizes = self._get_var_sizes([1, 2])
        
        if var_sizes.get(1) is None or var_sizes.get(2) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_1_size = var_sizes[1]
        var_num_2_size = var_sizes[2]
                # 判断是否存在缓冲区溢出
        if (var_num_2_size + 1 > var_num_1_size):
            return True
        return False


    def handle_memmove(self):
        """
        处理 memmove 函数的特定逻辑。
        """
        # 这里可以加入特定的判断逻辑，比如判断源缓冲区大小是否大于目标缓冲区
        var_sizes = self._get_var_sizes([1, 3])
        
        if var_sizes.get(1) is None or var_sizes.get(3) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_1_size = var_sizes[1]
        var_num_3_size = var_sizes[3]
                # 判断是否存在缓冲区溢出
        if (var_num_3_size > var_num_1_size):
            return True
        return False    
        return "memmove处理逻辑"

    def handle_strcat(self):
        """
        处理 strcat 函数的特定逻辑。
        """
        var_sizes = self._get_var_sizes([1, 2])
        
        if var_sizes.get(1) is None or var_sizes.get(2) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_1_size = var_sizes[1]
        var_num_2_size = var_sizes[2]
                # 判断是否存在缓冲区溢出
        if (var_num_2_size > var_num_1_size):
            return True
        return False  

    def handle_sprintf(self):
        """
        处理 sprintf 函数的特定逻辑。
        """
        return "sprintf处理逻辑"

    def handle_gets(self):
        """
        处理 gets 函数的特定逻辑。
        """
        return "gets处理逻辑"
    
    def handle_fgets(self):
        """
        处理 fgets 函数的特定逻辑。
        """
        # 这里可以加入特定的判断逻辑，比如判断源缓冲区大小是否大于目标缓冲区
        var_sizes = self._get_var_sizes([1, 2])
        
        if var_sizes.get(1) is None or var_sizes.get(2) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_1_size = var_sizes[1]
        var_num_2_size = var_sizes[2]
                # 判断是否存在缓冲区溢出
        if (var_num_2_size > var_num_1_size):
            return True
        return False        

    def handle_wcscpy(self):
        """
        处理 wcscpy 函数的特定逻辑。
        """
        var_sizes = self._get_var_sizes([1, 2])
        
        if var_sizes.get(1) is None or var_sizes.get(2) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_1_size = var_sizes[1]
        var_num_2_size = var_sizes[2]
                # 判断是否存在缓冲区溢出
        if (var_num_2_size > var_num_1_size):
            return True
        return False  

    def handle_wcscat(self):
        """
        处理 wcscat 函数的特定逻辑。
        """
        return "wcscat处理逻辑"
    
    def handle_recv(self):
        """
        处理 recv 函数的特定逻辑。
        """
        var_sizes = self._get_var_sizes([2, 3])
        
        if var_sizes.get(2) is None or var_sizes.get(3) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_2_size = var_sizes[2]
        var_num_3_size = var_sizes[3]
                # 判断是否存在缓冲区溢出
        if (var_num_2_size > var_num_3_size):
            return True
        return False  
    
    def handle_strncpy(self):
        """
        处理 strncpy 函数的特定逻辑。
        """
        var_sizes = self._get_var_sizes([1, 3])
        
        if var_sizes.get(1) is None or var_sizes.get(3) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_1_size = var_sizes[1]
        var_num_3_size = var_sizes[3]
                # 判断是否存在缓冲区溢出
        if (var_num_1_size < var_num_3_size):
            return True
        return False
    
    def handle_memset(self):
        """
        处理 memset 函数的特定逻辑。
        """
        var_sizes = self._get_var_sizes([1, 3])
        
        if var_sizes.get(1) is None or var_sizes.get(3) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_1_size = var_sizes[1]
        var_num_3_size = var_sizes[3]
                # 判断是否存在缓冲区溢出
        if (var_num_1_size < var_num_3_size):
            return True
        return False    
    
    def handle_strncat(self):
        """
        处理 strncat 函数的特定逻辑。
        """
        var_sizes = self._get_var_sizes([1, 3])
        
        if var_sizes.get(1) is None or var_sizes.get(3) is None:
            print("无法获取所有必要的变量信息")
            return False
        
        var_num_1_size = var_sizes[1]
        var_num_3_size = var_sizes[3]
                # 判断是否存在缓冲区溢出
        if (var_num_1_size < var_num_3_size):
            return True
        return False          

    def _get_var_sizes(self, var_nums):
        """
        根据给定的 var_nums 获取对应的 var_size。
        """
        var_sizes = {}
        for var_info in self.variable_info.values():
            var_num = var_info['var_num']
            if var_num in var_nums:
                var_size = int(var_info['var_size'])
                # print(var_size)
                var_sizes[var_num] = var_size
        return var_sizes

def check_for_buffer_overflow(function_name, variable_info):
    """
    根据function_name和variable_info判断是否存在缓冲区溢出漏洞。
    
    参数:
    function_name: 函数名称（如 'memcpy'）。
    variable_info: 字典，包含变量的信息（如 var_num、var_size）。
    
    返回:
    如果存在缓冲区溢出漏洞，返回提示信息；否则返回 None。
    """
    # print(function_name)
    # 判断 function_name 是否为 'memcpy'
    handler = SensitiveFunctionsHandler(variable_info)
    
    if function_name in sensitive_functions1:
        # 调用对应的处理函数
        handler_method = getattr(handler, f'handle_{function_name}', None)
        if handler_method:
            return handler_method()
        else:
            return "没有定义该函数的处理逻辑"
    
    return None

# def process_sensitive_functions(json_file_path):
#     """
#     处理从JSON文件中读取的敏感函数，并对调用地址进行反汇编和反编译分析。
#     """
#     # 读取 JSON 文件中的敏感函数和调用地址
#     sensitive_funcs = read_sensitive_functions_from_json(json_file_path)
    
#     for func in sensitive_funcs:
#         function_name = func["name"]
#         call_address = func["call_address"]
        
#         # 如果函数名在敏感函数列表中
#         if function_name in sensitive_functions1:
#             # 减去偏移量 0x400000
#             adjusted_address = int(call_address, 16) - 0x400000
            
#             # 获取反编译或反汇编代码行
#             print(f"Processing call to '{function_name}' at adjusted address {hex(adjusted_address)}...")
#             decompiled_line = get_decompiled_line_from_address(adjusted_address)
#             if decompiled_line:
#                 print(f"Decompiled line: {decompiled_line}")
#             else:
#                 print(f"Disassembled code: {idc.GetDisasm(adjusted_address)}")
            
#             params = extract_variables_from_call(decompiled_line)
#             # print(params)
#             variable_info = get_variable_type_and_size(adjusted_address , params)
            
#             for var, info in variable_info.items():
#                 if info['var_type']  == 'pointer':
#                     analyze_pointer_size(adjusted_address , info)
#                 print(f"{var}: {info['var_size']}, {info['var_type']}, {info['var_num']}")
                
#             result = check_for_buffer_overflow(function_name, variable_info)
#             if result:
#                 print("存在缓冲区溢出漏洞")
#             else:
#                 print("没有缓冲区溢出漏洞")

def process_sensitive_functions(json_file_path):
    """
    处理指定 JSON 文件中的敏感函数：
    - 对每个敏感函数调用地址进行分析
    - 在原 JSON 条目中新增 param_info 和 Vulnerable 字段
    - 出错时跳过该函数，继续下一个
    """
    data = read_sensitive_functions_from_json(json_file_path)
    # 对 direct_sensitive_func 和 inline_sensitive_func 两个列表都做相同处理
    for list_key, name_key in [
        ("direct_sensitive_func", "name"),
        ("inline_sensitive_func", "inline_func_name_predict")
    ]:
        if list_key not in data:
            continue

        for func in data[list_key]:
            try:
                function_name = func.get(name_key)
                call_address_str = func.get("call_address")
                if not function_name or not call_address_str:
                    continue

                # 只有在敏感函数列表中的才处理
                if function_name not in sensitive_functions1:
                    continue

                # 计算调整后的地址（减去偏移量 0x400000）
                adjusted_addr = int(call_address_str, 16) - 0x400000
                print(f"[+] 处理函数 {function_name} 地址 0x{adjusted_addr:x}")

                # 反编译或反汇编获取调用行
                decompiled_line = get_decompiled_line_from_address(adjusted_addr)
                if not decompiled_line:
                    decompiled_line = idc.GetDisasm(adjusted_addr)

                # 提取调用参数
                params = extract_variables_from_call(decompiled_line)
                # 获取各变量类型/大小信息
                variable_info = get_variable_type_and_size(adjusted_addr, params)
                for var, info in variable_info.items():
                    print(variable_info)

                # 对指针类型做额外分析
                for var, info in variable_info.items():
                    if info.get('var_type') == 'pointer':
                        analyze_pointer_size(adjusted_addr, info)

                # 检测是否存在缓冲区溢出漏洞
                is_vuln = check_for_buffer_overflow(function_name, variable_info)

                # 更新 JSON 条目
                func["param_info"] = variable_info
                func["Vulnerable"] = bool(is_vuln)
                status = "存在漏洞" if is_vuln else "无漏洞"
                print(f"    -> 分析结果：{status}")

                # 每处理完一个函数，就写回文件
                write_sensitive_functions_to_json(json_file_path, data)
                print("    [*] 已写回 JSON 文件")


            except Exception as e:
                # 捕获所有异常，打印错误并继续下一个函数
                print(f"[!] 处理 {function_name}@{call_address_str} 时出错：{e}")
                write_sensitive_functions_to_json(json_file_path, data)
                continue
            
    print("[*] 全部处理完成，结果已写回 JSON 文件。")

def main():
    # print("[*] 开始分析对 strcpy 的调用...")
    # # 运行敏感函数调用查找
    # json_file_path = r'D:\Users\Administrator\Desktop\study\project\dataset\test\json\bad\CWE121_Stack_Based_Buffer_Overflow__CWE129_connect_socket_02\result.json'
    json_file_path = idc.ARGV[1]
    print(json_file_path)
    process_sensitive_functions(json_file_path)
    print("[*] 分析结束。")

if __name__ == '__main__':
    idaapi.autoWait()
    main()
    idc.qexit(0)